from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from src.agentic.config.VectorStore import VectorStore
from langchain_core.output_parsers import StrOutputParser
from langchain_community.utils.math import cosine_similarity

logger = getLogger()

class RoutingSemantic:

    def __init__(self, llm_model, embed_model, collection_prefix, document_sources):
        self.llm_model = llm_model
        self.embed_model = embed_model
        self.collection_prefix = collection_prefix
        self.document_sources = document_sources

    def invoke(self, query):
        logger.info(f"RoutingSemantic invoke query: {query}")

        datasource_list = []
        datasource_descs = []
        for source in self.document_sources:
            datasource_list.append(self.collection_prefix + source.library_number)
            datasource_descs.append(source.document_summary)
        logger.info(f"RoutingLogic invoke datasource_list: {datasource_list}")

        datasource_embedding = self.embed_model.embed_documents(datasource_descs)
        query_embedding = self.embed_model.embed_query(query)
        similarity = cosine_similarity([query_embedding], datasource_embedding)
        similarity_datasource = datasource_list[similarity.argmax()]
        logger.info(f"RoutingSemantic invoke similarity_datasource: {similarity_datasource}")

        vector_store = VectorStore().new_vector_store(self.embed_model, similarity_datasource)
        retriever = vector_store.as_retriever(search_kwargs = { "k": 3 })
        retrieve_docs = retriever.invoke(query)

        template = """
            请基于以下上下文内容回答问题：
            {context}
            
            问题：{question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        chain_result = chain.invoke({ "context": retrieve_docs, "question": query })
        logger.info(f"RoutingSemantic invoke chain_result len: {len(chain_result)}")
        retrieve_doc = [doc.page_content for doc in retrieve_docs]
        return { "retrieve_docs": retrieve_doc, "chain_result": chain_result }
